#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdbool.h>

#define AINV_SIZE 16    // AINV矩阵维度
#define CUBIC_DEGREE 4 // 三次多项式维度（0-3次）

// ========================= 全局常量：AINV矩阵（与Python一致） =========================
static const double AINV[AINV_SIZE][AINV_SIZE] = {
    {1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {-3, 3, 0, 0, -2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {2, -2, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0},
    {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0},
    {0, 0, 0, 0, 0, 0, 0, 0, -3, 3, 0, 0, -2, -1, 0, 0},
    {0, 0, 0, 0, 0, 0, 0, 0, 2, -2, 0, 0, 1, 1, 0, 0},
    {-3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0, 0, 0, 0, 0},
    {0, 0, 0, 0, -3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0},
    {9, -9, -9, 9, 6, 3, -6, -3, 6, -6, 3, -3, 4, 2, 2, 1},
    {-6, 6, 6, -6, -3, -3, 3, 3, -4, 4, -2, 2, -2, -2, -1, -1},
    {2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0},
    {0, 0, 0, 0, 2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0},
    {-6, 6, 6, -6, -4, -2, 4, 2, -3, 3, -3, 3, -2, -1, -2, -1},
    {4, -4, -4, 4, 2, 2, -2, -2, 2, -2, 2, -2, 1, 1, 1, 1}
};

// ========================= 双变量单段函数结构体（复刻BiVariablesFunction） =========================
typedef struct {
    double matrix[CUBIC_DEGREE][CUBIC_DEGREE]; // 4x4系数矩阵
    double x_domain[2];                        // x区间 [x0, x1]
    double y_domain[2];                        // y区间 [y0, y1]
} BiVariablesFunction;

/**
 * 创建双变量单段函数
 * @param matrix 4x4系数矩阵
 * @param x0 x区间左边界
 * @param x1 x区间右边界
 * @param y0 y区间左边界
 * @param y1 y区间右边界
 * @return 双变量函数实例
 */
BiVariablesFunction* biv_func_create(double matrix[CUBIC_DEGREE][CUBIC_DEGREE], 
                                     double x0, double x1, double y0, double y1) {
    BiVariablesFunction* func = (BiVariablesFunction*)malloc(sizeof(BiVariablesFunction));
    if (!func) {
        perror("malloc failed for BiVariablesFunction");
        return NULL;
    }
    // 拷贝系数矩阵
    for (int i = 0; i < CUBIC_DEGREE; i++) {
        memcpy(func->matrix[i], matrix[i], CUBIC_DEGREE * sizeof(double));
    }
    func->x_domain[0] = x0;
    func->x_domain[1] = x1;
    func->y_domain[0] = y0;
    func->y_domain[1] = y1;
    return func;
}

/**
 * 释放双变量单段函数内存
 * @param func 函数实例
 */
void biv_func_destroy(BiVariablesFunction* func) {
    if (func) free(func);
}

/**
 * 双变量单段函数求值（复刻__call__）
 * @param func 函数实例
 * @param x 自变量x
 * @param y 自变量y
 * @return 函数值
 */
double biv_func_evaluate(const BiVariablesFunction* func, double x, double y) {
    if (!func) return NAN;

    // 归一化x、y到[0,1]区间
    double x_norm = (x - func->x_domain[0]) / (func->x_domain[1] - func->x_domain[0]);
    double y_norm = (y - func->y_domain[0]) / (func->y_domain[1] - func->y_domain[0]);

    // 计算y的幂向量 [1, y, y², y³]
    double y_power[CUBIC_DEGREE] = {1.0, y_norm, y_norm*y_norm, y_norm*y_norm*y_norm};
    // 计算x的幂向量 [1, x, x², x³]
    double x_power[CUBIC_DEGREE] = {1.0, x_norm, x_norm*x_norm, x_norm*x_norm*x_norm};

    double result = 0.0;
    for (int i = 0; i < CUBIC_DEGREE; i++) {
        double r = 0.0;
        // 系数矩阵行 × y幂向量
        for (int j = 0; j < CUBIC_DEGREE; j++) {
            r += func->matrix[i][j] * y_power[j];
        }
        // 结果 × x幂向量
        result += r * x_power[i];
    }

    return result;
}

// ========================= 辅助函数：查找区间索引（复刻index_of） =========================
/**
 * 查找值所属区间的左索引
 * @param values 升序数组
 * @param len 数组长度
 * @param x 目标值
 * @return 区间左索引（-1表示超出范围）
 */
int index_of(const double* values, int len, double x) {
    for (int i = 0; i < len - 1; i++) {
        if (x >= values[i] - 1e-10 && x <= values[i + 1] + 1e-10) { // 浮点容错
            return i;
        }
    }
    return -1; // 超出范围
}

// ========================= 双变量分段函数结构体（复刻BiVariablesFunctions） =========================
typedef struct {
    BiVariablesFunction*** functions; // 二维分段函数数组 [x_idx][y_idx]
    int x_count;                      // x节点数量
    int y_count;                      // y节点数量
    double* x_values;                 // x节点数组
    double* y_values;                 // y节点数组
} BiVariablesFunctions;

/**
 * 创建双变量分段函数
 * @param funcs 二维函数数组
 * @param x_vals x节点数组
 * @param x_len x节点数量
 * @param y_vals y节点数组
 * @param y_len y节点数量
 * @return 分段函数实例
 */
BiVariablesFunctions* biv_funcs_create(BiVariablesFunction*** funcs, 
                                       double* x_vals, int x_len,
                                       double* y_vals, int y_len) {
    BiVariablesFunctions* funcs_obj = (BiVariablesFunctions*)malloc(sizeof(BiVariablesFunctions));
    if (!funcs_obj) {
        perror("malloc failed for BiVariablesFunctions");
        return NULL;
    }
    funcs_obj->functions = funcs;
    funcs_obj->x_count = x_len;
    funcs_obj->y_count = y_len;
    
    // 拷贝节点数组
    funcs_obj->x_values = (double*)malloc(x_len * sizeof(double));
    funcs_obj->y_values = (double*)malloc(y_len * sizeof(double));
    if (!funcs_obj->x_values || !funcs_obj->y_values) {
        perror("malloc failed for x/y values");
        free(funcs_obj->x_values);
        free(funcs_obj->y_values);
        free(funcs_obj);
        return NULL;
    }
    memcpy(funcs_obj->x_values, x_vals, x_len * sizeof(double));
    memcpy(funcs_obj->y_values, y_vals, y_len * sizeof(double));
    
    return funcs_obj;
}

/**
 * 释放双变量分段函数内存
 * @param funcs_obj 分段函数实例
 */
void biv_funcs_destroy(BiVariablesFunctions* funcs_obj) {
    if (!funcs_obj) return;

    // 释放二维函数数组
    int n_x = funcs_obj->x_count - 1;
    int n_y = funcs_obj->y_count - 1;
    for (int i = 0; i < n_x; i++) {
        for (int j = 0; j < n_y; j++) {
            biv_func_destroy(funcs_obj->functions[i][j]);
        }
        free(funcs_obj->functions[i]);
    }
    free(funcs_obj->functions);

    // 释放节点数组
    free(funcs_obj->x_values);
    free(funcs_obj->y_values);
    free(funcs_obj);
}

/**
 * 双变量分段函数求值（复刻__call__）
 * @param funcs_obj 分段函数实例
 * @param x 自变量x
 * @param y 自变量y
 * @return 函数值（NAN表示超出范围）
 */
double biv_funcs_evaluate(const BiVariablesFunctions* funcs_obj, double x, double y) {
    if (!funcs_obj) return NAN;

    // 查找x、y所属区间索引
    int idx_x = index_of(funcs_obj->x_values, funcs_obj->x_count, x);
    int idx_y = index_of(funcs_obj->y_values, funcs_obj->y_count, y);

    if (idx_x < 0 || idx_y < 0) {
        fprintf(stderr, "x=%.2f or y=%.2f out of range\n", x, y);
        return NAN;
    }

    // 调用对应分段函数求值
    return biv_func_evaluate(funcs_obj->functions[idx_x][idx_y], x, y);
}

// ========================= 双立方插值核心函数 =========================
/**
 * 计算偏导数（复刻diffs函数）
 * @param x_vals x节点数组
 * @param x_len x节点数量
 * @param y_vals y节点数组
 * @param y_len y节点数量
 * @param f_vals f(x,y)值矩阵 [x_idx][y_idx]
 * @param dx 输出：x偏导数矩阵
 * @param dy 输出：y偏导数矩阵
 * @param dx_dy 输出：混合偏导数矩阵
 */
void diffs(const double* x_vals, int x_len,
           const double* y_vals, int y_len,
           const double** f_vals,
           double** dx, double** dy, double** dx_dy) {
    // 初始化偏导数矩阵为0
    for (int i = 0; i < x_len; i++) {
        memset(dx[i], 0, y_len * sizeof(double));
        memset(dy[i], 0, y_len * sizeof(double));
        memset(dx_dy[i], 0, y_len * sizeof(double));
    }

    // 计算内部点的偏导数（边界点保持0）
    for (int i = 1; i < x_len - 1; i++) {
        double diff_x = x_vals[i + 1] - x_vals[i - 1];
        for (int j = 1; j < y_len - 1; j++) {
            double diff_y = y_vals[j + 1] - y_vals[j - 1];
            
            // y偏导数
            dy[i][j] = (f_vals[i][j + 1] - f_vals[i][j - 1]) / diff_y;
            // x偏导数
            dx[i][j] = (f_vals[i + 1][j] - f_vals[i - 1][j]) / diff_x;
            // 混合偏导数（注意：Python原代码分母是diff_x*diff_x，保持一致）
            dx_dy[i][j] = (f_vals[i + 1][j + 1] + f_vals[i - 1][j - 1]
                          - f_vals[i - 1][j + 1] - f_vals[i + 1][j - 1]) / (diff_x * diff_x);
        }
    }
}

/**
 * 计算系数矩阵（复刻coefficients函数）
 * @param beta 16维beta向量
 * @param matrix 输出：4x4系数矩阵
 */
void coefficients(const double* beta, double matrix[CUBIC_DEGREE][CUBIC_DEGREE]) {
    double coeff[AINV_SIZE] = {0.0};

    // 计算coeff = AINV × beta
    for (int i = 0; i < AINV_SIZE; i++) {
        double r = 0.0;
        for (int j = 0; j < AINV_SIZE; j++) {
            r += AINV[i][j] * beta[j];
        }
        coeff[i] = r;
    }

    // 转换为4x4矩阵（i%4=列，i//4=行）
    for (int i = 0; i < AINV_SIZE; i++) {
        int col = i % CUBIC_DEGREE;
        int row = i / CUBIC_DEGREE;
        matrix[row][col] = coeff[i];
    }
}

/**
 * 双立方插值主函数（复刻interpolate）
 * @param x_vals x节点数组
 * @param x_len x节点数量
 * @param y_vals y节点数组
 * @param y_len y节点数量
 * @param f_vals f(x,y)值矩阵 [x_idx][y_idx]
 * @return 分段双变量函数实例
 */
BiVariablesFunctions* bicubic_interpolate(const double* x_vals, int x_len,
                                          const double* y_vals, int y_len,
                                          const double** f_vals) {
    if (x_len < 2 || y_len < 2) {
        fprintf(stderr, "x/y节点数量至少为2\n");
        return NULL;
    }

    int n = x_len - 1; // 分段数

    // 分配偏导数矩阵内存
    double** dx = (double**)malloc(x_len * sizeof(double*));
    double** dy = (double**)malloc(x_len * sizeof(double*));
    double** dx_dy = (double**)malloc(x_len * sizeof(double*));
    for (int i = 0; i < x_len; i++) {
        dx[i] = (double*)malloc(y_len * sizeof(double));
        dy[i] = (double*)malloc(y_len * sizeof(double));
        dx_dy[i] = (double*)malloc(y_len * sizeof(double));
        if (!dx[i] || !dy[i] || !dx_dy[i]) {
            perror("malloc failed for diff matrices");
            // 回滚内存
            for (int j = 0; j <= i; j++) {
                free(dx[j]); free(dy[j]); free(dx_dy[j]);
            }
            free(dx); free(dy); free(dx_dy);
            return NULL;
        }
    }

    // 计算偏导数
    diffs(x_vals, x_len, y_vals, y_len, f_vals, dx, dy, dx_dy);

    // 分配分段函数数组 [n][n]
    BiVariablesFunction*** functions = (BiVariablesFunction***)malloc(n * sizeof(BiVariablesFunction**));
    if (!functions) {
        perror("malloc failed for functions array");
        goto cleanup_diffs;
    }
    for (int i = 0; i < n; i++) {
        functions[i] = (BiVariablesFunction**)malloc(n * sizeof(BiVariablesFunction*));
        if (!functions[i]) {
            perror("malloc failed for functions row");
            // 回滚
            for (int j = 0; j < i; j++) free(functions[j]);
            free(functions);
            goto cleanup_diffs;
        }
    }

    // 构建每个分段的双变量函数
    for (int i = 0; i < n; i++) {
        double diff_x = x_vals[i + 1] - x_vals[i];
        for (int j = 0; j < n; j++) {
            double diff_y = y_vals[j + 1] - y_vals[j];

            // 构建16维beta向量
            double beta[AINV_SIZE] = {
                f_vals[i][j], f_vals[i+1][j], f_vals[i][j+1], f_vals[i+1][j+1],
                dx[i][j] * diff_x, dx[i+1][j] * diff_x, dx[i][j] * diff_x, dx[i][j+1] * diff_x,
                dy[i][j] * diff_x, dy[i+1][j] * diff_x, dy[i][j] * diff_x, dy[i][j+1] * diff_x,
                dx_dy[i][j] * diff_y, dx_dy[i+1][j] * diff_y, dx_dy[i][j] * diff_y, dx_dy[i][j+1] * diff_y
            };

            // 计算4x4系数矩阵
            double matrix[CUBIC_DEGREE][CUBIC_DEGREE];
            coefficients(beta, matrix);

            // 创建双变量单段函数
            functions[i][j] = biv_func_create(matrix, 
                                             x_vals[i], x_vals[i+1],
                                             y_vals[j], y_vals[j+1]);
            if (!functions[i][j]) {
                // 回滚已创建的函数
                for (int ii = 0; ii <= i; ii++) {
                    for (int jj = 0; jj < (ii == i ? j : n); jj++) {
                        biv_func_destroy(functions[ii][jj]);
                    }
                    free(functions[ii]);
                }
                free(functions);
                goto cleanup_diffs;
            }
        }
    }

    // 创建分段函数实例
    BiVariablesFunctions* result = biv_funcs_create(functions, x_vals, x_len, y_vals, y_len);

    // 释放临时内存
    for (int i = 0; i < n; i++) free(functions[i]);
    free(functions);

cleanup_diffs:
    // 释放偏导数矩阵
    for (int i = 0; i < x_len; i++) {
        free(dx[i]); free(dy[i]); free(dx_dy[i]);
    }
    free(dx); free(dy); free(dx_dy);

    return result;
}

// ========================= 测试函数 =========================
void test_1() {
    printf("=== Test 1 ===\n");
    // 测试数据
    double x_values[] = {0, 5, 10};
    double y_values[] = {0, 5, 10};
    double* f_values[] = {
        (double[]){5, 3, 5},
        (double[]){3, 10, 3},
        (double[]){5, 3, 5}
    };
    int x_len = sizeof(x_values)/sizeof(double);
    int y_len = sizeof(y_values)/sizeof(double);

    // 插值计算
    BiVariablesFunctions* funcs = bicubic_interpolate(x_values, x_len, y_values, y_len, (const double**)f_values);
    if (!funcs) return;

    // 输出节点值
    for (int i = 0; i < x_len; i++) {
        for (int j = 0; j < y_len; j++) {
            double x = x_values[i];
            double y = y_values[j];
            double val = biv_funcs_evaluate(funcs, x, y);
            printf("f(%.0f,%.0f)= %.1f\n", x, y, val);
        }
    }

    // 输出插值点
    double val = biv_funcs_evaluate(funcs, 3, 3);
    printf("f(3,3)= %.10f\n\n", val);

    // 释放内存
    biv_funcs_destroy(funcs);
}

void test_2() {
    printf("=== Test 2 ===\n");
    // 测试数据
    double x_values[] = {0, 5, 10};
    double y_values[] = {0, 4, 8};
    double* f_values[] = {
        (double[]){5, 3, 5},
        (double[]){3, 10, 3},
        (double[]){5, 3, 5}
    };
    int x_len = sizeof(x_values)/sizeof(double);
    int y_len = sizeof(y_values)/sizeof(double);

    // 插值计算
    BiVariablesFunctions* funcs = bicubic_interpolate(x_values, x_len, y_values, y_len, (const double**)f_values);
    if (!funcs) return;

    // 输出节点值
    for (int i = 0; i < x_len; i++) {
        for (int j = 0; j < y_len; j++) {
            double x = x_values[i];
            double y = y_values[j];
            double val = biv_funcs_evaluate(funcs, x, y);
            printf("f(%.0f,%.0f)= %.1f\n", x, y, val);
        }
    }

    // 输出插值点
    double val = biv_funcs_evaluate(funcs, 3, 3);
    printf("f(3,3)= %.5f\n\n", val);

    // 释放内存
    biv_funcs_destroy(funcs);
}

int main() {
    test_1();
    test_2();
    return 0;
}